from langchain_core.tools import tool
from web_apps import app, db
from web_apps.datamodel.db_models import DataModel as DataModelObj
from utils.etl_utils import get_reader_model
from web_apps.llm.utils import get_llm, extract_code


@tool
def get_reader(name: str) -> object:
    '''
    根据数据模型名称获取对应数据读取对象reader
    '''
    with app.app_context():
        obj = db.session.query(DataModelObj).filter(DataModelObj.name.like(f'%{name}%')).first()
        if obj is not None:
            flag, reader = get_reader_model({'model_id': obj.id})
            if flag:
                return reader
            else:
                return f"获取数据读取对象失败：{reader}"
        return f"获取数据读取对象失败：未找到对应数据"


@tool
def gen_reader_info_prompt(reader: object) -> str:
    '''
    获取数据读取对象reader使用说明和元信息
    '''
    with app.app_context():
        return reader.get_info_prompt()


@tool
def generate_code(reader_info_prompt: str, prompt: str) -> str:
    '''
    生成数据处理代码
    '''
    result_example_prompt = '{ "type": "string", "value": "100" } or { "type": "dataframe", "value": pd.DataFrame({...}) } or { "type": "html", "value": line.render_embed() }'
    prompt = f"""
我有一个数据读取对象reader，对象信息为：
{reader_info_prompt}

Update this initial code:
```python
# TODO: import the required dependencies

# Write code here

# Declare result var: 
type (possible values "string", "dataframe", "html"). Example: {result_example_prompt}

```

### QUERY

{prompt}

Variable `reader` is already declared.

At the end, declare "result" variable as a dictionary of type and value.

If you are asked to plot a chart, use "pyecharts" for charts, use the render_embed() function to return the corresponding html type and the html content value.

Generate python code and return full updated code:
生成代码前请使用中文解释大致逻辑
"""
    llm = get_llm()
    llm_result = llm.invoke(prompt).content
    code = extract_code(llm_result)
    return code


@tool
def execute_query_code(reader: object, code: str) -> any:
    """
    运行数据处理代码获取结果
    Execute the python code generated by LLMs to answer the question
    about the input dataframe. Run the code in the current context and return the
    result.
    Args:
        code (str): Python code to execute.
    Returns:
        Any: The result of the code execution. The type of the result depends
            on the generated code.
    """
    try:
        environment = {'reader': reader}
        exec(code, environment)
        if "result" not in environment:
            raise ValueError("No result returned")
        else:
            result = environment['result']
            return result
    except Exception as e:
        raise e
